LlamaIndexを完全に理解するチュートリアル その2:テキスト分割のカスタマイズ
こんちには。
データアナリティクス事業本部 インテグレーション部 機械学習チームの中村です。
「LlamaIndexを完全に理解するチュートリアル その2」では、テキスト分割のカスタマイズを取り扱います。
本記事で使用する用語は以下のその1で説明していますので、そちらも参照ください。
LlamaIndexを完全に理解するチュートリアル |
---|
その1:処理の概念や流れを理解する基礎編(v0.7.9対応) |
その2:テキスト分割のカスタマイズ |
・本記事の内容はその1のv0.7.9版の記事を投稿後、v0.7.9で動作するように修正しています
本記事の内容
LlamaIndexは与えられたテキストをインデックス化しますが、LLMとのやり取りではテキストの長さ(トークン数)の上限という制約があります。
具体的に、ChatGPTの裏側で使用されるgpt-3.5-turbo
などは、トークン数が4096個が上限となり、日本語の文字数でもおおよそ同じ数が上限となります。
この制約は、「入力(プロンプト)」と「応答」双方の合計値に対するものですので、文脈を与えてチャットに応答される場合、与えられる文脈には制限が出てきます。
そのため、LlamaIndexでは、テキストをある単位でチャンク分割する処理を行います。
通常デフォルトで使用する場合は、チャンク分割が暗黙的に行われていますので、今回はこのカスタマイズ方法を見ていきます。
環境準備
その1と同様の方法で準備します。
使用したバージョン情報は以下となります。
- Python : 3.10.11
- langchain : 0.0.234
- llama-index : 0.7.9
- openai : 0.27.8
サンプルコード
ベースのサンプルは以下とします。こちらを実行しておきます。
from llama_index import SimpleDirectoryReader from llama_index import ListIndex documents = SimpleDirectoryReader(input_dir="./data").load_data() list_index = ListIndex.from_documents(documents) query_engine = list_index.as_query_engine()
ノード分割の状況を可視化
上記を実行後に、以下でノードの分割状況を可視化できます。
for doc_id, node in list_index.storage_context.docstore.docs.items(): node_dict = node.__dict__ print(f'{doc_id=}, len={len(node_dict["text"])}, start={node_dict["start_char_idx"]}, end={node_dict["end_char_idx"]}')
doc_id='a8ecf0bb-e664-4ba3-a77d-ec9b7d591b3e', len=903, start=0, end=903 doc_id='b4e33e3f-5a82-4076-acb8-2af54add4c46', len=899, start=904, end=1803 doc_id='c9b3c44b-c429-4911-bdc7-4146f2dc0de3', len=942, start=1798, end=2740 doc_id='a280f12a-2c37-4ae7-93c9-42b675453422', len=245, start=2747, end=2992 doc_id='e0296a5b-a7b3-4885-8e44-94f7bc845fb9', len=735, start=0, end=735 doc_id='0688f16e-1627-416a-bba9-2dd526f09dad', len=795, start=736, end=1531 doc_id='b18bcc33-2f3f-4c66-9c2c-96be8c306e4a', len=832, start=1532, end=2364 doc_id='ac537ef0-a364-4866-99c2-5efcabc58a6b', len=616, start=2365, end=2981
こちらがデフォルトの動作です。
カスタマイズするためのコード修正
チャンク分割はServiceContextのNodeParserのTextSplitterが担っています。
実際にはTextSplitterはLangChainで定義されているクラスで、LlamaIndexはTextSplitterのサブクラスとしてTokenTextSplitterなどを定義しています。
カスタマイズするの場合は以下のように表にTextSplitterを出します。
from llama_index import SimpleDirectoryReader from llama_index import Document from llama_index import GPTListIndex from llama_index import ServiceContext from llama_index.node_parser import SimpleNodeParser from llama_index.langchain_helpers.text_splitter import TokenTextSplitter from llama_index.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE import tiktoken documents = SimpleDirectoryReader(input_dir="./data").load_data() text_splitter = TokenTextSplitter(separator=" ", chunk_size=DEFAULT_CHUNK_SIZE , chunk_overlap=DEFAULT_CHUNK_OVERLAP , tokenizer=tiktoken.get_encoding("gpt2").encode) node_parser = SimpleNodeParser(text_splitter=text_splitter) service_context = ServiceContext.from_defaults( node_parser=node_parser ) list_index = GPTListIndex.from_documents(documents , service_context=service_context)
最初はデフォルトの動作と同じパラメータでTextSplitterを作成しました。
再度可視化してみましょう。
for doc_id, node in list_index.storage_context.docstore.docs.items(): node_dict = node.__dict__ print(f'{doc_id=}, len={len(node_dict["text"])}, start={node_dict["start_char_idx"]}, end={node_dict["end_char_idx"]}')
doc_id='a8ecf0bb-e664-4ba3-a77d-ec9b7d591b3e', len=903, start=0, end=903 doc_id='b4e33e3f-5a82-4076-acb8-2af54add4c46', len=899, start=904, end=1803 doc_id='c9b3c44b-c429-4911-bdc7-4146f2dc0de3', len=942, start=1798, end=2740 doc_id='a280f12a-2c37-4ae7-93c9-42b675453422', len=245, start=2747, end=2992 doc_id='e0296a5b-a7b3-4885-8e44-94f7bc845fb9', len=735, start=0, end=735 doc_id='0688f16e-1627-416a-bba9-2dd526f09dad', len=795, start=736, end=1531 doc_id='b18bcc33-2f3f-4c66-9c2c-96be8c306e4a', len=832, start=1532, end=2364 doc_id='ac537ef0-a364-4866-99c2-5efcabc58a6b', len=616, start=2365, end=2981 doc_id='f1622414-e837-460a-8321-34eccfc61993', len=903, start=0, end=903 doc_id='69cee1eb-f457-45b1-9855-d6b259554dab', len=899, start=904, end=1803 doc_id='f1a4f4eb-9799-4389-bac8-e3aae306f945', len=942, start=1798, end=2740 doc_id='c9fe8605-f5e4-404d-9ac8-ca74892cb852', len=245, start=2747, end=2992 doc_id='e1c891b3-cdef-4227-8620-a6e317b89f64', len=735, start=0, end=735 doc_id='35b4cb29-e1bc-470b-b877-8f3026bf6440', len=795, start=736, end=1531 doc_id='256979b7-0bca-4048-b63b-32cf724a5faf', len=832, start=1532, end=2364 doc_id='10e81af7-e023-461d-bcc5-1db9a0b9db99', len=616, start=2365, end=2981
同じ結果となることが確認できました。
カスタマイズの詳細
テキスト分割は以下の部分に集約されます。
text_splitter = TokenTextSplitter(separator=" ", chunk_size=DEFAULT_CHUNK_SIZE , chunk_overlap=DEFAULT_CHUNK_OVERLAP , tokenizer=tiktoken.get_encoding("gpt2").encode)
各パラメータの意味は以下です。
separator
: 区切り文字。ここの指定文字がチャンク分割の切れ目になるようチャンクを作ります。chunk_size
: チャンクサイズ。このチャンクサイズ以下となるようチャンクを作ります。chunk_overlap
: チャンクのオーバーラップ。後ろのチャンクに前のチャンクの末尾を付け加えることで、チャンク分割による文脈の断裂を緩和します。tokenizer
: tokenizerのencode処理。チャンクサイズをトークン数でカウントするために使用されます。
separatorについて
separator
はデフォルトは半角空白ですが、今回のサンプルデータでは句読点の後ろにたまたま半角空白が入っているため、句読点でうまく区切れているようです。
空白が無い場合で、句読点できちんと区切ってほしい場合は句読点を設定する必要がありそうです。
またseparator
はstr型に対するsplitと同じ動作で処理しますので、複数文字を設定すると、それをひと固まりの区切り文字として処理しますので、注意が必要です。
"hoge fuga. fuga hoge.".split(". ")
['hoge fuga', 'fuga hoge.']
ですので、ORで区切り文字を構成するには前処理で工夫するか、TokenTextSplitterでre.split
を使うようにサブクラスを再定義するかなどの必要がありそうですのでご注意ください。
tokenizerについて
チャンクサイズはトークン数単位でカウントするようですので、それにこのtokenizerが使用されます。
可視化をトークン数でやり直すと以下のようになります。
import tiktoken for doc_id, node in list_index.storage_context.docstore.docs.items(): node_dict = node.__dict__ print(f'{doc_id=}, len={len(tiktoken.get_encoding("gpt2").encode(node_dict["text"]))}, start={node_dict["start_char_idx"]}, end={node_dict["end_char_idx"]}')
doc_id='f1ada414-fa64-4ca9-8c5f-7ac8b098c231', len=1013, start=0, end=903 doc_id='3b2bf3ae-c82d-4c1d-9f09-c3d5f77da1e1', len=1019, start=744, end=1651 doc_id='da959ec2-1932-49c7-b984-744306af0e54', len=964, start=1698, end=2567 doc_id='534ed4a4-b1a3-4237-b56b-3397c96f8707', len=807, start=2532, end=3272 doc_id='3bed0a84-4369-49f2-9d88-be757798a9fd', len=959, start=0, end=735 doc_id='b743ff61-09d1-4754-acb8-191045da82c3', len=997, start=621, end=1413 doc_id='64ef19cd-75fd-4cc4-88af-36b3be715c0d', len=1004, start=1374, end=2208 doc_id='204f050a-4410-4b84-a753-b10e5d65fc5a', len=991, start=2208, end=3032 doc_id='edb612e7-4c47-4c34-ae51-3b415eee2641', len=444, start=3069, end=3430
このようにチャンクサイズがDEFAULT_CHUNK_SIZE=1024以下のトークン数に収まっていることが分かります。
またデフォルトではgpt2
のtokenizerが使用されている点も注意が必要です。
厳密には、tokenizerが後段のLLMやEmbeddingに合ったものを使用するのが理想的です。以下が代表的なものとなります。
text-davinci-003
:p50k_base
gpt-3.5-turbo
:cl100k_base
gpt-4
:cl100k_base
text-embedding-ada-002
:cl100k_base
近年のモデルを使う場合、gpt2
よりもcl100k_base
の方が適切かなという印象です。
詳細はtiktokenのコードを参照ください。
カスタマイズしてみた
以下のカスタマイズをした例を作ります。
tokenizer
にcl100k_base
を使用chunk_size
を512
に設定chunk_overlap
を100
に設定seperator
は。
の句点と空白を設定
from llama_index import SimpleDirectoryReader from llama_index import Document from llama_index import GPTListIndex from llama_index import ServiceContext from llama_index.node_parser import SimpleNodeParser from llama_index.langchain_helpers.text_splitter import TokenTextSplitter from llama_index.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE import tiktoken documents = SimpleDirectoryReader(input_dir="./data").load_data() text_splitter = TokenTextSplitter(separator="。 ", chunk_size=512 , chunk_overlap=100 , tokenizer=tiktoken.get_encoding("cl100k_base").encode) node_parser = SimpleNodeParser(text_splitter=text_splitter) service_context = ServiceContext.from_defaults( node_parser=node_parser ) list_index = GPTListIndex.from_documents(documents , service_context=service_context)
念のため可視化してチェックします。
for doc_id, node in list_index.storage_context.docstore.docs.items(): node_dict = node.__dict__ print(f'{doc_id=}, len={len(node_dict["text"])}, start={node_dict["start_char_idx"]}, end={node_dict["end_char_idx"]}')
doc_id='5e5e6602-dbd3-4c73-9dd4-043d5b6151c5', len=554, start=0, end=554 doc_id='004d6b65-6d96-4655-9540-edd643e64858', len=570, start=458, end=1028 doc_id='d2efac46-5175-42df-8973-579c5585f69a', len=332, start=1049, end=1381 doc_id='036a7e5b-55e2-40b7-a74f-be3af525b1a7', len=467, start=1377, end=1844 doc_id='ce97efc6-9a8c-4220-8264-db1a906af1dc', len=480, start=1825, end=2305 doc_id='2a0cf7ae-ecbe-41c3-a6a4-97416c879fe1', len=554, start=2408, end=2962 doc_id='7b3cb425-bb3d-46fd-a1b4-decbbbfb2d34', len=414, start=2944, end=3358 doc_id='3bfb585f-7073-43e0-8bf0-e19e9a030705', len=492, start=0, end=492 doc_id='9de00bef-f127-42ee-aca6-56190de39c58', len=388, start=432, end=820 doc_id='c7367d75-9eb5-4a47-8cb4-68b35483a999', len=506, start=869, end=1375 doc_id='8ef855eb-faf8-462e-93f0-0e05bcdff09a', len=523, start=1307, end=1830 doc_id='1a2c124d-b64f-4009-a8ef-b89434e197e1', len=526, start=1822, end=2348 doc_id='cf265404-547d-4d6f-b944-07a9734c78cf', len=473, start=2397, end=2870 doc_id='8c2a40c0-6ba4-4039-92ae-d9ed65453cba', len=397, start=2900, end=3297
意図通りになっていることが分かりました。
まとめ
いかがでしたでしょうか。
こちらをベースにカスタマイズにぜひチャレンジして見てください。
本記事が、今後LlamaIndexをお使いになられる方の参考になれば幸いです。